import os
directory = os.fsencode('Data/aligned')
files = []
for sub1 in os.listdir(directory):
    sub1_name = os.fsdecode(sub1)
    if '.fits'in sub1_name:
        files.append('Data/aligned'+"/"+sub1_name)

import os
from tqdm import tqdm
from astropy.io import fits
from astropy.time import TimeDelta, Time
import torch
import numpy as np
from astropy.io import fits
import matplotlib.pyplot as plt
from astropy.visualization import (ImageNormalize, ZScaleInterval, LogStretch, MinMaxInterval, AsinhStretch)
from astropy.wcs import WCS
from scipy.ndimage import shift
from astropy.coordinates import SkyCoord

total_data = {}
pass_counter = 0
MAGZPL_LIST = []
for i in tqdm(range(len(files))):

    # Open the FITS file
    filename = files[i]
    hdul = fits.open(filename)    
    # Access the primary HDU (Header Data Unit)
    primary_hdu = hdul[0]
    
    # Get the data and header
    # try:
    data = primary_hdu.data
    header = primary_hdu.header

    total_elements = data.size
    # Count the number of NaN values
    nan_count = np.isnan(data).sum()
    # Calculate the percentage of NaN values
    nan_percentage = (nan_count / total_elements) * 100
    if nan_percentage < 20:
        time = header['SHUTOPEN']
        MAGZPL_LIST.append(header['MAGZP'])
        time = Time(time.replace('T',' '), format='iso', scale='utc')
        total_data[time] = torch.tensor(data.astype(np.float32))
    else:
        pass_counter += 1
        pass
    hdul.close()
print(pass_counter, 'skipped out of ', len(files)) 
MAGZPL = np.max(MAGZPL_LIST)
print(MAGZPL)

100%|█████████████████████████████████████████| 939/939 [00:16<00:00, 57.27it/s]

671 skipped out of  939
26.416529

sorted_keys = sorted(total_data.keys())
earliest_frame = sorted_keys[0]

# Apply normalization
norm = ImageNormalize(stretch=AsinhStretch(a=0.0001), interval=ZScaleInterval())  
data = total_data[earliest_frame]
# norm = ImageNormalize(total_data[earliest_frame], interval=ZScaleInterval())

# Plot the image with normalization
plt.figure(figsize=(10, 8))
plt.imshow(data, cmap='gray_r', origin='lower', norm=norm)
plt.colorbar(label='Pixel Value')
plt.title('FITS Image with Normalization')
plt.xlabel('RA Pixel')
plt.ylabel('Dec Pixel')
plt.show()

from astropy.io import fits
import matplotlib.pyplot as plt
from astropy.visualization import (ImageNormalize, ZScaleInterval, LogStretch, MinMaxInterval, AsinhStretch)
import numpy as np
import matplotlib.pyplot as plt
from celluloid import Camera
from IPython.display import HTML

norm = ImageNormalize(stretch=AsinhStretch(a=0.0001), interval=ZScaleInterval())  

%matplotlib inline
# Set up the figure and axes
fig, ax = plt.subplots(figsize=(8,8))

# Initialize Celluloid Camera
camera = Camera(fig)


data = total_data[earliest_frame]
count = 0
# Animate the sine wave with a changing phase
for key in sorted_keys:
    ax.imshow(total_data[key], cmap='gray_r', origin='lower', norm=norm)  # Plot the sine wave
    camera.snap()  # Capture the frame
    if count > 100:
        break
    else:
        count += 1

# Create the animation
animation = camera.animate(interval=100)  # 100ms between frames

# Display the animation
HTML(animation.to_jshtml())  # Render as JS HTML

import numpy as np
import matplotlib.pyplot as plt

def generate_2d_gaussian(grid_size, amplitude=1, x0=0, y0=0, sigma_x=1, sigma_y=1):
    Nx, Ny = grid_size
    x = np.linspace(-Nx // 2, Nx // 2, Nx)
    y = np.linspace(-Ny // 2, Ny // 2, Ny)
    X, Y = np.meshgrid(x, y)
    gaussian = amplitude * np.exp(-(((X - x0) ** 2) / (2 * sigma_x ** 2) +
                                    ((Y - y0) ** 2) / (2 * sigma_y ** 2)))
    return gaussian

ra_0, dec_0 = 0,0 # center of the frame
dx, dy = 1,  2 # we shift up and to the right these are in units of pixel/day for the projected velocities
MAG = 22.51 # magnitude of the injected star

total_data = {}
pass_counter = 0


for i in tqdm(range(len(files))):

    # Open the FITS file
    filename = files[i]
    hdul = fits.open(filename)    
    # Access the primary HDU (Header Data Unit)
    primary_hdu = hdul[0]
    
    # Get the data and header
    data = primary_hdu.data
    header = primary_hdu.header
    # ====== make sure to keep only images that are NOT all NANs
    total_elements = data.size
    # Count the number of NaN values
    nan_count = np.isnan(data).sum()
    # Calculate the percentage of NaN values
    nan_percentage = (nan_count / total_elements) * 100
    if nan_percentage < 20:
        time = header['SHUTOPEN']
        time = Time(time.replace('T',' '), format='iso', scale='utc')
        # ========== step 1: remove saturated pixels ==========
        saturate = header['SATURATE']
        data = np.where(data > saturate/2, np.nan, data) 
        # =========  step 2: subtract the median (sky subtraction) =======
        median = np.nanmedian(data)
        data = data - median
        # =========  step 3: scale flux =============
        MAGZP = header['MAGZP']
        scale = 10**(0.4*(MAGZPL - MAGZP))
        data = data * scale
        # =========== step 3.5: compute limiting Magnitude ======
        see = header['SEEING']
        pix = header['PIXSCALE']
        var = np.nanmedian(abs(data))
        lmt = -2.5*np.log10(3.0*np.sqrt(np.pi*see*see/pix/pix)*var*1.48) + MAGZPL
        # =========== step 4: drop gaussian fake =========
        grid_size = data.shape[::-1] # Grid size (Nx, Ny)
        dt = (time - earliest_frame).to_value('day') 
        x0, y0 = 0 + dt * dx, 0 + dt * dy  # shift the gaussian
        
        FWHM = header['SEEING']
        sigma = FWHM/2.3548
        sigma_x, sigma_y = sigma, sigma  # Standard deviations
        amplitude = 10**(-0.4*(MAG-MAGZPL)) *(1/(2*np.pi*sigma**2))
        gaussian = generate_2d_gaussian(grid_size, amplitude, x0, y0, sigma_x, sigma_y)
        print(gaussian.max())
        data = np.array(data) + gaussian
       
        header = primary_hdu.header
        
        total_data[time] = torch.tensor(data.astype(np.float32))
    else:
        pass_counter += 1
        pass
    hdul.close()
print(pass_counter, 'skipped out of ', len(files))  
# 8.940820307431759

  1%|▏                                          | 5/939 [00:00<01:09, 13.51it/s]

0.0
0.0

  1%|▎                                          | 7/939 [00:00<02:26,  6.36it/s]

8.85885037414935

  1%|▎                                          | 8/939 [00:01<02:51,  5.42it/s]

0.0

  1%|▍                                         | 10/939 [00:01<02:43,  5.68it/s]

8.469222277449635

  1%|▌                                         | 13/939 [00:01<02:16,  6.77it/s]

0.0

sorted_keys = sorted(total_data.keys())
earliest_frame = sorted_keys[0]

from astropy.io import fits
import matplotlib.pyplot as plt
from astropy.visualization import (ImageNormalize, ZScaleInterval, LogStretch, MinMaxInterval, AsinhStretch)
import numpy as np
import matplotlib.pyplot as plt
from celluloid import Camera
from IPython.display import HTML

norm = ImageNormalize(stretch=AsinhStretch(a=0.0001), interval=ZScaleInterval())  

%matplotlib inline
# Set up the figure and axes
fig, ax = plt.subplots(figsize=(10,10))

# Initialize Celluloid Camera
camera = Camera(fig)


data = total_data[earliest_frame]
count = 0
# Animate the sine wave with a changing phase
for key in sorted_keys:
    ax.imshow(total_data[key], cmap='gray_r', origin='lower', norm=norm)  # Plot the sine wave
    camera.snap()  # Capture the frame
    if count > 100:
        break
    else:
        count += 1

# Create the animation
animation = camera.animate(interval=100)  # 100ms between frames

# Display the animation
HTML(animation.to_jshtml())  # Render as JS HTML

import torch
import torch.nn as nn
import torch.nn.functional as F
from time import time

class Shift_Stack(nn.Module):
    def __init__(self, frames, times):
        super(Shift_Stack, self).__init__()
        self.frames = self.frame_subtraction(frames).cuda()[:,None,:,:]
        self.delta_times = torch.zeros_like(times).cuda()
        self.delta_times[1:] = times[1:] - times[0]
        self.B = self.frames.shape[0] # predefine batch sizes (number of frames)
        self.grid = F.affine_grid(
            torch.eye(2, 3).unsqueeze(0).repeat(self.B, 1, 1),  # Identity matrix for rotation and scaling
            size=self.frames.size(),  # Output size
            align_corners=False
        ).cuda()
        
    # @torch.jit.script
    def frame_subtraction(self, frame):
        """
        Sequential frame subtractions
        """
        frame = frame - frame[0,:,:]
        return frame
        
    def decimate_32_to_16(self, data_float32):
        """
        Decimate data from float32 to float16 while preserving the range
        """
        # Define a scaling factor to bring values within float16 range
        scale_factor = data_float32.abs().max() / 6.55e4  # Approx. max for float16
        scale_factor = max(scale_factor, 1.0)  # Ensure we only scale down
        
        # Scale, convert to float16, and scale back
        data_scaled = data_float32 / scale_factor
        data_float16 = data_scaled.to(dtype=torch.float16)
        data_preserved = data_float16 * scale_factor
        return data_preserved
        
    def average_nonzeros(self, data, axis=0):
        """
        This averages only the nonzero values
        """
        # Mask to identify non-zero values
        non_zero_sum = torch.sum(data, dim=axis)
        non_zero_count = torch.count_nonzero(data, dim=axis)
        
        # Avoid division by zero
        non_zero_count = torch.where(non_zero_count == 0, torch.tensor(1.0, device=data.device), non_zero_count)
        
        # Calculate average of non-zero elements
        non_zero_avg = non_zero_sum / non_zero_count
        return non_zero_avg

    def shift_images(self, images, shifts):
        """
        Shifts a batch of images by specified (x, y) values for each image.
        
        images (torch.Tensor): A batch of images of shape (B, C, H, W),
                               where B is batch size, C is number of channels,
                               H is height, and W is width.
        shifts (torch.Tensor): A tensor of shape (B, 2), where each row
                               represents the (B, x, y) shift for the corresponding image.
        """
        B, C, H, W = images.shape
        shifts_x, shifts_y = shifts[:, 0], shifts[:, 1]
        # Create a grid for affine transformations
        start = time()
        grid = self.grid.clone()
        # Adjust the grid by the shifts
        grid[..., 0] += 2 * shifts_x.view(-1, 1, 1) / W  # Normalize shift in x direction
        grid[..., 1] += 2 * shifts_y.view(-1, 1, 1) / H  # Normalize shift in y direction
        start = time()
        # Perform the grid sampling to shift the images
        shifted_images = F.grid_sample(
            images, grid, mode='bilinear', padding_mode='zeros', align_corners=False
        ).cuda()
        return shifted_images

    def forward(self, dxdy):
        """
        Preprocessing handing before calling shift_images
        """
        dx, dy = dxdy[:,0], dxdy[:,1]
        shifts_x = torch.multiply(dx , self.delta_times) # pixel shift in x
        shifts_y = torch.multiply(dy , self.delta_times) # pixel shift in y
        shifts = torch.stack([shifts_x, shifts_y] , axis = 1)
        frames = self.frames
        batched_shift = self.shift_images(frames, shifts)
        batched_shift = batched_shift[:,0,:,: ] # meant for future multi-filters... [curr. index zero for 1 filter]
        # stacked = torch.mean(batched_shift, axis = 0) # FASTER MEAN Calculations
        stacked = torch.nanmedian(batched_shift, axis = 0)
        # stacked = self.average_nonzeros(batched_shift, axis = 0)
        return stacked

frames = []
frame_times = []
adjusts = []
for i, key in enumerate(sorted_keys):
    if i < 150:
        frames.append(total_data[key])
        frame_times.append(key.to_value('jd'))
frames = torch.stack(frames)
frame_times = torch.tensor(frame_times)
print(frames.shape, frame_times.shape)

torch.Size([150, 3080, 3072]) torch.Size([150])

stacking_job = Shift_Stack(frames, frame_times)

shifts = torch.ones((150, 2)).cuda()
shifts[:, 0] = dx
shifts[:, 1] = dy
print(shifts.shape)

torch.Size([150, 2])

start = time()
result = stacking_job(shifts)
print(time()-start, "runtime [s]")
start = time()
result = stacking_job(shifts)
print(time()-start, "runtime [s]")
start = time()
result = stacking_job(shifts)
print(time()-start, "runtime [s]")
start = time()
result = stacking_job(shifts)
print(time()-start, "runtime [s]")
start = time()
result = stacking_job(shifts)
print(time()-start, "runtime [s]")
start = time()
result = stacking_job(shifts)
print(time()-start, "runtime [s]")
start = time()
result = stacking_job(shifts)
print(time()-start, "runtime [s]")

0.12030720710754395 runtime [s]
0.0007495880126953125 runtime [s]
0.00016570091247558594 runtime [s]
0.00015783309936523438 runtime [s]
0.00015497207641601562 runtime [s]
0.0001575946807861328 runtime [s]
0.0001533031463623047 runtime [s]

norm = ImageNormalize(stretch=AsinhStretch(a=100), interval=ZScaleInterval())  
# 
plt.figure(figsize=(10, 10))
plt.imshow(result.values.cpu().numpy(), cmap='gray_r', origin='lower', vmin = -10, vmax=10 )
# plt.imshow(result.values.cpu().numpy(), cmap='gray_r', origin='lower', vmin = result.values.cpu().numpy().min(), vmax=result.values.cpu().numpy().max() )
plt.colorbar(label='Pixel Value')
plt.title('FITS Image with Normalization')
plt.xlabel('RA Pixel')
plt.ylabel('Dec Pixel')
plt.show()

plt.figure(figsize=(10, 10))
plt.imshow(result.values.cpu().numpy()[1500:1600,1500:1600], cmap='gray_r', origin='lower', vmin = result.values.cpu().numpy()[1500:1600,1500:1600].min(), vmax=result.values.cpu().numpy()[1500:1600,1500:1600].max() )
plt.colorbar(label='Pixel Value')
plt.title('FITS Image with Normalization')
plt.xlabel('RA Pixel')
plt.ylabel('Dec Pixel')
plt.show()

%matplotlib inline

array = result.values.cpu().numpy()
flat_index = np.argmax(array)

# Convert flat index to 2D indices
row, col = np.unravel_index(flat_index, data.shape)

print("Maximum value:", array[row, col])
print("Index of maximum value:", (row, col))

Maximum value: 10.842533
Index of maximum value: (np.int64(62), np.int64(2148))

Pytorch Shift Stack Demo¶

Loading Data + Preprocessing¶

Add fakes¶

Watch the injections!¶

Shift and Add¶